/**
 * Creat with IntelliJ IDEA
 * Description:
 * User:WZW
 * Date:2022-08-29
 * Time:10:50
 */
public class AVLTree {

    static class TreeNode {
        public int val;
        public int bf;// 平衡因子
        public TreeNode left;// 左子树的引用
        public TreeNode right;// 右子树的引用
        public TreeNode parent;// 父亲节点的引用

        public TreeNode(int val) {
            this.val = val;
        }
    }

    public TreeNode root;// 根节点

    public boolean insert(int val) {
        TreeNode node = new TreeNode(val);
        if (root == null) {
            root = node;
            return true;
        }

        TreeNode parent = null;
        TreeNode cur = root;
        // 查找插入的位置
        while (cur != null) {
            if (cur.val < val) {
                parent = cur;
                cur = cur.right;
            } else if (cur.val == val) {
                return false;
            } else {
                parent = cur;
                cur = cur.left;
            }
        }
        if (parent.val < val) {
            parent.right = node;
        } else {
            parent.left = node;
        }
        node.parent = parent;
        cur = node;
        // 调节平衡因子
        while (parent != null) {
            // 先看cur 是 parent 的做还是右 -> 平衡因子的改变
            if (cur == parent.right) {
                parent.bf++; // 如果是右树，说明右树高度增加，bf++
            }else {
                parent.bf--; //如果是左树，说明左树高度增加，bf--
            }

            // 检查当前的平衡因子
            if (parent.bf == 0) {
                // 说明已经平衡了
            }else if (parent.bf == 1 || parent.bf == -1) {
                // 需要继续向上修改平衡因子
                cur = parent;
                parent = cur.parent;
            }else {
                if (parent.bf == 2) {
                    //右树高，需要降低右树的高度
                    if (cur.bf == 1) {
                        //左旋
                        rotateLeft(parent);
                    }else {
                        //cur.bf == -1
                        rotateRL(parent);
                    }
                }else {// 左树高
                    //parent.bf == -2
                    if (cur.bf == -1) {
                        //右旋
                        rotateRight(parent);
                    }else {
                        //cur.bf == 1
                        rotateLR(parent);
                    }
                }
                // 上述代码走完就平衡了
                break;
            }
        }
        return true;
    }

    /*
    右左双旋
     */
    private void rotateRL(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;
        rotateRight(parent.right);
        rotateLeft(parent);
        if (bf == 1) {
            parent.bf = -1;
            subR.bf = 0;
            subRL.bf = 0;
        } else if (bf == -1) {
            parent.bf = 0;
            subR.bf = 1;
            subRL.bf = 0;
        }
    }

    /*
    左右双旋：先左旋，后右旋
     */
    private void rotateLR(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf;

        rotateLeft(parent.left);
        rotateRight(parent);

        if (bf == -1) {
            subL.bf = 0;
            subLR.bf = 0;
            parent.bf = 1;
        }else if (bf == 1){
            subL.bf = -1;
            subLR.bf = 0;
            parent.bf = 0;
        }

    }

    /*
    左单旋
     */
    public void rotateLeft(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;

        parent.right = subRL;
        if (subRL != null) {
            subRL.parent = parent;
        }
        subR.left = parent;

        if (parent == root) {
            root = subR;
            root.parent = null;
        } else {
            TreeNode pParent = parent.parent;
            if (pParent.left == parent) {
                pParent.left = subR;
            }else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        parent.parent = subR;
        parent.bf = 0;
        subR.bf = 0;
    }

    /*
    右单旋
     */
    public void rotateRight(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;

        parent.left = subLR;
        subL.right = parent;
        if (subLR != null) {
            subLR.parent = parent;
        }

        TreeNode pParent = parent.parent;
        parent.parent = subL;
        // 检查当前是不是根节点
        if (parent == root) {
            root = subL;
            root.parent = null;
        } else {
            // 不是根节点，判断这颗子树是左子树还是右子树
            if (pParent.left == parent) {
                pParent.left = subL;
            }else {
                pParent.right = subL;
            }
            subL.parent = parent.parent;
        }
        parent.bf = 0;
        subL.bf = 0;
        return;
    }

    public boolean isBalanced(TreeNode root) {
        if (root == null) return true;
        int leftH = height(root.left);
        int rightH = height(root.right);

        if (rightH - leftH != root.bf) {
            System.out.println("节点" + root.val + "平衡因子异常");
            return false;
        }

        return Math.abs(leftH - rightH) <= 1
                && isBalanced(root.left)
                && isBalanced(root.right);
    }

    private int height(TreeNode root) {
        if (root == null) return 0;
        int leftH = height(root.left);
        int rightH = height(root.right);
        return leftH > rightH ? leftH + 1 : rightH + 1;
    }

}
